package com.mou.aiagent.rag;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 自定义文档加载器，加载多个md文档
 * @return allDocuments
 */
@Component
@Slf4j
public class LoveAppDocumentLoader {
    // 定义一个ResourcePatternResolver类型的私有常量 读取多个文件
    private  final ResourcePatternResolver resourcePatternResolver;
    // 构造函数，接收一个ResourcePatternResolver类型的参数
    public LoveAppDocumentLoader(ResourcePatternResolver resourcePatternResolver) {
        // 将参数赋值给私有常量
        this.resourcePatternResolver = resourcePatternResolver;
    }
    // 定义一个方法，接收一个String类型的参数，返回一个DocumentLoader类型的对象
    public List<Document>loadDocuments() {
        // 定义集合接收所有文档
        List<Document> allDocuments = new ArrayList<>();
        //加载多篇文档
        try {
            // 获取所有Markdown文件
            Resource[] resources = resourcePatternResolver.getResources("classpath:document/*.md");
            for (Resource resource : resources) {
                // 获取文件名
                String filename = resource.getFilename();
                //获取文件的长度作为标签
                String status = filename.substring(filename.length() - 6, filename.length() - 3);
                // Markdown加载器
                MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                        .withHorizontalRuleCreateDocument(true)
                        .withIncludeCodeBlock(false)
                        .withIncludeBlockquote(false)

                // 添加额外的元数据，键为"filename"，值为filename
                        .withAdditionalMetadata("filename", filename )
                        .withAdditionalMetadata("status",status)
                        .build();
                // 加载文档
                // 创建一个MarkdownDocumentReader对象，用于读取Markdown文档
                MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader(resource, config);
               //将文档添加到集合中
               allDocuments.addAll(markdownDocumentReader.get());
                // 将文档添加到集合中
            }
        } catch (IOException e) {
          log.error("Error loading documents: {}", e.getMessage());
        }
        return allDocuments;
    }

}
